/** @file
AppleImageConversion protocol

Copyright (C) 2018 savvas. All rights reserved.<BR>
Portions copyright (C) 2016 slice. All rights reserved.<BR>
Portions copyright (C) 2018 vit9696. All rights reserved.<BR>

This program and the accompanying materials
are licensed and made available under the terms and conditions of the BSD License
which accompanies this distribution.  The full text of the license may be found at
http://opensource.org/licenses/bsd-license.php

THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
**/

#include <Library/BaseMemoryLib.h>
#include <Library/DebugLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/OcAppleImageConversionLib.h>
#include <Library/OcGuardLib.h>
#include <Library/OcMiscLib.h>
#include <Library/OcPngLib.h>
#include <Library/UefiBootServicesTableLib.h>

STATIC
CONST UINT8
mPngHeader[] = { 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A };

STATIC
EFI_STATUS
EFIAPI
RecognizeImageData (
  IN VOID   *ImageBuffer,
  IN UINTN  ImageSize
  )
{
  if (ImageBuffer == NULL
    || ImageSize < sizeof (mPngHeader)
    || CompareMem (ImageBuffer, mPngHeader, sizeof (mPngHeader)) != 0) {
    return EFI_INVALID_PARAMETER;
  }

  return EFI_SUCCESS;
}

STATIC
EFI_STATUS
EFIAPI
GetImageDims (
  IN  VOID    *ImageBuffer,
  IN  UINTN   ImageSize,
  OUT UINT32  *ImageWidth,
  OUT UINT32  *ImageHeight
  )
{
  EFI_STATUS  Status;

  Status = GetPngDims (ImageBuffer, ImageSize, ImageWidth, ImageHeight);

  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_INFO, "Failed to obtain image dimensions for image\n"));
  }

  return Status;
}

STATIC
EFI_STATUS
EFIAPI
DecodeImageData (
  IN  VOID           *ImageBuffer,
  IN  UINTN          ImageSize,
  OUT EFI_UGA_PIXEL  **RawImageData,
  OUT UINTN          *RawImageDataSize
  )
{
  EFI_STATUS      Status;
  UINTN           Index;
  UINTN           PixelCount;
  EFI_UGA_PIXEL   *PixelWalker;
  UINT32          Width;
  UINT32          Height;
  UINT8           TmpChannel;

  STATIC_ASSERT (sizeof (EFI_UGA_PIXEL) == sizeof (UINT32), "Unsupported pixel size");
  STATIC_ASSERT (OFFSET_OF (EFI_UGA_PIXEL, Blue)     == 0,  "Unsupported pixel format");
  STATIC_ASSERT (OFFSET_OF (EFI_UGA_PIXEL, Green)    == 1,  "Unsupported pixel format");
  STATIC_ASSERT (OFFSET_OF (EFI_UGA_PIXEL, Red)      == 2,  "Unsupported pixel format");
  STATIC_ASSERT (OFFSET_OF (EFI_UGA_PIXEL, Reserved) == 3,  "Unsupported pixel format");

  if (RawImageData == NULL || RawImageDataSize == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  Status = DecodePng (
    ImageBuffer,
    ImageSize,
    (VOID **) RawImageData,
    &Width,
    &Height,
    NULL
    );

  if (EFI_ERROR (Status)) {
    return EFI_UNSUPPORTED;
  }

  PixelCount  = (UINTN) Width * Height;
  PixelWalker = *RawImageData;

  for (Index = 0; Index < PixelCount; ++Index) {
    TmpChannel            = PixelWalker->Blue;
    PixelWalker->Blue     = PixelWalker->Red;
    PixelWalker->Red      = TmpChannel;
    PixelWalker->Reserved = 0xFF - PixelWalker->Reserved;
    ++PixelWalker;
  }
  
  return EFI_SUCCESS;
}

STATIC
EFI_STATUS
EFIAPI
GetImageDimsVersion (
  IN VOID     *Buffer,
  IN  UINTN   BufferSize,
  IN  UINTN   Version,
  OUT UINT32  *Width,
  OUT UINT32  *Height
  )
{
  if (Buffer == NULL
    || BufferSize == 0
    || Version == 0
    || Width == NULL
    || Height == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  if (Version > APPLE_IMAGE_CONVERSION_PROTOCOL_INTERFACE_V1) {
    return EFI_UNSUPPORTED;
  }

  return GetImageDims (Buffer, BufferSize, Width, Height);
}

STATIC
EFI_STATUS
EFIAPI
DecodeImageDataVersion (
  IN  VOID           *Buffer,
  IN  UINTN          BufferSize,
  IN  UINTN          Version,
  OUT EFI_UGA_PIXEL  **RawImageData,
  OUT UINTN          *RawImageDataSize
  )
{
  if (Buffer == NULL
    || BufferSize == 0
    || Version == 0
    || RawImageData == NULL
    || RawImageDataSize == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  if (Version > APPLE_IMAGE_CONVERSION_PROTOCOL_INTERFACE_V1) {
    return EFI_UNSUPPORTED;
  }

  return DecodeImageData (Buffer, BufferSize, RawImageData, RawImageDataSize);
}

//
// Image codec protocol instance.
//
STATIC APPLE_IMAGE_CONVERSION_PROTOCOL mAppleImageConversion = {
  APPLE_IMAGE_CONVERSION_PROTOCOL_REVISION,
  APPLE_IMAGE_CONVERSION_PROTOCOL_ANY_EXTENSION,
  RecognizeImageData,
  GetImageDims,
  DecodeImageData,
  GetImageDimsVersion,
  DecodeImageDataVersion
};

APPLE_IMAGE_CONVERSION_PROTOCOL *
OcAppleImageConversionInstallProtocol (
  IN BOOLEAN  Reinstall
  )
{
  EFI_STATUS                       Status;
  APPLE_IMAGE_CONVERSION_PROTOCOL  *AppleImageConversionInterface;
  EFI_HANDLE                       NewHandle;

  if (Reinstall) {
    Status = UninstallAllProtocolInstances (&gAppleImageConversionProtocolGuid);
    if (EFI_ERROR (Status)) {
      DEBUG ((DEBUG_ERROR, "OCIC: Uninstall failed: %r\n", Status));
      return NULL;
    }
  } else {
    Status = gBS->LocateProtocol (
      &gAppleImageConversionProtocolGuid,
      NULL,
      (VOID **) &AppleImageConversionInterface
      );
    if (!EFI_ERROR (Status)) {
      return AppleImageConversionInterface;
    }
  }

  NewHandle = NULL;
  Status = gBS->InstallMultipleProtocolInterfaces (
    &NewHandle,
    &gAppleImageConversionProtocolGuid,
    &mAppleImageConversion,
    NULL
    );

  if (EFI_ERROR (Status)) {
    return NULL;
  }

  return &mAppleImageConversion;
}
